import os
import time

import rosbag
import numpy as np
import sensor_msgs.point_cloud2 as pc2
import open3d as o3d
import cv2
import math

def pc_tf(pc, RT):
    # extrinsic 中旋转的表达形式为旋转矩阵
    # pc shape (n , 3)
    # extrinsic shape (4, 4)
    xyz = pc[:, :3]
    xyz1 = np.concatenate((xyz, np.ones((len(xyz), 1))), axis=1)
    pc2 = pc.copy()
    pc2[:, :3] = ((RT @ xyz1.T).T)[:, :3]
    return pc2


def npy2ply(npy):
    ply = o3d.geometry.PointCloud()
    ply.points = o3d.utility.Vector3dVector(npy[:, :3])
    # density = npy[:, 3]
    colors = [[0.5, 0.7, 0.5]] * npy.shape[0]
    ply.colors = o3d.utility.Vector3dVector(colors)
    return ply


def bbox3d2corners(bboxes):
    '''
    bboxes: shape=(n, 7)
    return: shape=(n, 8, 3)
           ^ z   x            6 ------ 5
           |   /             / |     / |
           |  /             2 -|---- 1 |
    y      | /              |  |     | |
    <------|o               | 7 -----| 4
                            |/   o   |/
                            3 ------ 0
    x: front, y: left, z: top
    '''
    centers, dims, angles = bboxes[:, :3], bboxes[:, 3:6], bboxes[:, 6]

    # 1.generate bbox corner coordinates, clockwise from minimal point
    bboxes_corners = np.array([[-0.5, -0.5, 0], [-0.5, -0.5, 1.0], [-0.5, 0.5, 1.0], [-0.5, 0.5, 0.0],
                               [0.5, -0.5, 0], [0.5, -0.5, 1.0], [0.5, 0.5, 1.0], [0.5, 0.5, 0.0]],
                              dtype=np.float32)
    bboxes_corners = bboxes_corners[None, :, :] * dims[:, None, :]  # (1, 8, 3) * (n, 1, 3) -> (n, 8, 3)

    # 2. rotate around z axis
    rot_sin, rot_cos = np.sin(angles), np.cos(angles)
    # in fact, -angle
    rot_mat = np.array([[rot_cos, rot_sin, np.zeros_like(rot_cos)],
                        [-rot_sin, rot_cos, np.zeros_like(rot_cos)],
                        [np.zeros_like(rot_cos), np.zeros_like(rot_cos), np.ones_like(rot_cos)]],
                       dtype=np.float32)  # (3, 3, n)
    rot_mat = np.transpose(rot_mat, (2, 1, 0))  # (n, 3, 3)
    bboxes_corners = bboxes_corners @ rot_mat  # (n, 8, 3)

    # 3. translate to centers
    bboxes_corners += centers[:, None, :]
    return bboxes_corners


def rotz(t):
    c = np.cos(t)
    s = np.sin(t)
    return np.array([[c, -s, 0], [s, c, 0], [0, 0, 1]])


def box2corner(box):
    x = box[0]
    y = box[1]
    z = box[2]
    l = box[3]  # dx
    w = box[4]  # dy
    h = box[5]  # dz
    yaw = box[6]
    Box = np.array(
        [
            [-l / 2, -l / 2, l / 2, l / 2, -l / 2, -l / 2, l / 2, l / 2],
            [w / 2, -w / 2, -w / 2, w / 2, w / 2, -w / 2, -w / 2, w / 2],
            [-h / 2, -h / 2, -h / 2, -h / 2, h / 2, h / 2, h / 2, h / 2],
        ]
    )
    # 先旋转再平移
    R = rotz(yaw)
    corners_3d = np.dot(R, Box)  # corners_3d: (3, 8)
    corners_3d[0, :] = corners_3d[0, :] + x
    corners_3d[1, :] = corners_3d[1, :] + y
    corners_3d[2, :] = corners_3d[2, :] + z
    return np.transpose(corners_3d)


def get_line_set(corners):
    lines = [
        [0, 1], [1, 2], [2, 3], [3, 0],
        [4, 5], [5, 6], [6, 7], [7, 4],
        [0, 4], [1, 5], [2, 6], [3, 7], ]
    colors = [[1, 0, 0] for i in range(len(lines))]
    line_set = o3d.geometry.LineSet()
    line_set.points = o3d.utility.Vector3dVector(corners)
    line_set.lines = o3d.utility.Vector2iVector(lines)
    line_set.colors = o3d.utility.Vector3dVector(colors)
    return line_set


def get_draw_box(pre_box_set):
    draw_boxes = []
    for box_p in pre_box_set:
        cx, cy, cz, sx, sy, sz, yaw = box_p[0], box_p[1], box_p[2], box_p[3], box_p[4], box_p[5], box_p[6]
        corner_box = box2corner([cx, cy, cz, sx, sy, sz, yaw])
        draw_box = get_line_set(corner_box)
        draw_boxes.append(draw_box)
    return draw_boxes

def vector_magnitude(vec):
    """
    Calculates a vector's magnitude.
    Args:
        - vec ():
    """
    magnitude = np.sqrt(np.sum(vec**2))
    return(magnitude)


def calculate_zy_rotation_for_arrow(vec):
    """
    Calculates the rotations required to go from the vector vec to the
    z axis vector of the original FOR. The first rotation that is
    calculated is over the z axis. This will leave the vector vec on the
    XZ plane. Then, the rotation over the y axis.

    Returns the angles of rotation over axis z and y required to
    get the vector vec into the same orientation as axis z
    of the original FOR

    Args:
        - vec ():
    """
    # Rotation over z axis of the FOR
    gamma = np.arctan(vec[1]/vec[0])
    Rz = np.array([[np.cos(gamma),-np.sin(gamma),0],
                   [np.sin(gamma),np.cos(gamma),0],
                   [0,0,1]])
    # Rotate vec to calculate next rotation
    vec = Rz.T@vec.reshape(-1,1)
    vec = vec.reshape(-1)
    # Rotation over y axis of the FOR
    beta = np.arctan(vec[0]/vec[2])
    Ry = np.array([[np.cos(beta),0,np.sin(beta)],
                   [0,1,0],
                   [-np.sin(beta),0,np.cos(beta)]])
    return(Rz, Ry)

def create_arrow(scale=10):
    """
    Create an arrow in for Open3D
    """
    cone_height = scale*0.2
    cylinder_height = scale*0.8
    cone_radius = scale/10
    cylinder_radius = scale/20
    mesh_frame = o3d.geometry.TriangleMesh.create_arrow(cone_radius=1,
        cone_height=cone_height,
        cylinder_radius=0.5,
        cylinder_height=cylinder_height)
    return(mesh_frame)

def get_arrow(origin=[0, 0, 0], end=None, vec=None):
    """
    Creates an arrow from an origin point to an end point,
    or create an arrow from a vector vec starting from origin.
    Args:
        - end (): End point. [x,y,z]
        - vec (): Vector. [i,j,k]
    """
    scale = 10
    Ry = Rz = np.eye(3)
    T = np.array([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]])
    T[:3, -1] = origin
    if end is not None:
        vec = np.array(end) - np.array(origin)
    elif vec is not None:
        vec = np.array(vec)
    if end is not None or vec is not None:
        scale = vector_magnitude(vec)
        Rz, Ry = calculate_zy_rotation_for_arrow(vec)
    mesh = create_arrow(scale)
    # Create the arrow
    mesh.rotate(Ry, center=np.array([0, 0, 0]))
    mesh.rotate(Rz, center=np.array([0, 0, 0]))
    mesh.translate(origin)
    return(mesh)

def get_arrows(pre_box_set):
    draw_arrows = []
    for box_p in pre_box_set:
        cx, cy, cz, sx, sy, sz, yaw = box_p[0], box_p[1], box_p[2], box_p[3], box_p[4], box_p[5], box_p[6]
        orin = [cx, cy, cz]
        end = [cx + 5 * math.cos(yaw), cy + 5 * math.sin(yaw), cz]
        draw_arrow = get_arrow(orin, end)
        draw_arrows.append(draw_arrow)
    return draw_arrows

class Vis():
    def __init__(self, pc_topics=["lidar", "radar"], box_topics=["boxes"], img_topics=["/camera71/compressed"]):
        self.vis = o3d.visualization.VisualizerWithKeyCallback()
        self.vis.create_window()
        self.vis.add_geometry(o3d.geometry.TriangleMesh.create_coordinate_frame(size=5, origin=[0, 0, 0]))
        self.datas = dict()
        for topic in pc_topics:
            self.datas[topic] = o3d.geometry.PointCloud()
            self.vis.add_geometry(self.datas[topic])

        lines_box = np.array([[0, 1], [1, 2], [0, 3], [2, 3], [4, 5], [4, 7], [5, 6], [6, 7],
                              [0, 4], [1, 5], [2, 6], [3, 7]])
        for i in range(0, 100):
            topic = "box" + str(i)
            self.datas[topic] = o3d.geometry.LineSet()
            # 将八个顶点连接次序的信息转换成o3d可以使用的数据类型
            self.datas[topic].lines = o3d.utility.Vector2iVector(lines_box)
            self.vis.add_geometry(self.datas[topic])

        # self.ctr = self.vis.get_view_control()
        self.to_reset = True

    def show_pc(self, points, topic, color=[0.3, 0.5, 0]):
        self.datas[topic].points = o3d.utility.Vector3dVector(points[:, :3])
        self.datas[topic].paint_uniform_color(color)
        self.vis.update_geometry(self.datas[topic])

        # self.vis.poll_events()
        # self.vis.update_renderer()
        #
        # if self.to_reset:
        #     self.vis.reset_view_point(True)
        #     self.to_reset = False

    def show_img(self, image):
        data = o3d.geometry.Image(image)
        self.vis.add_geometry(data)

    def show_box(self, bboxes, color=[[0.3, 0.5, 0]]):
        if len(bboxes.shape) == 2:
            bboxes = bbox3d2corners(bboxes)
        assert len(bboxes) < 100  # 笨方法显示box，否则不能调整视角
        for i in range(0, len(bboxes)):
            topic = "box" + str(i)
            # 设置每条线段的颜色
            self.datas[topic].colors = o3d.utility.Vector3dVector(color)
            self.datas[topic].points = o3d.utility.Vector3dVector(bboxes[i])
            # 将矩形框加入到窗口中
            self.vis.update_geometry(self.datas[topic])

    def show(self):
        self.vis.poll_events()
        self.vis.update_renderer()
        if self.to_reset:
            self.vis.reset_view_point(True)
            self.to_reset = False
